import xml.etree.ElementTree as ET
import os
import json
from tqdm import tqdm


# 解析labelimg xml
def labelimg(xml_path, txt_path):
    tree = ET.parse(xml_path)
    # img_name = tree.find("path").text.replace("/", "\\").split("\\")[-1]
    img_w = float(tree.find("size/width").text)
    img_h = float(tree.find("size/height").text)
    objs = tree.findall("object")

    # 直接打开写入，写完关闭
    w_txt = open(txt_path, "w")
    for obj in objs:
        clas_index = clas_name.index(obj.find("name").text)  # 类别索引
        # clas_index = 0  # 类别索引
        x1 = float(obj.find("bndbox/xmin").text)
        y1 = float(obj.find("bndbox/ymin").text)
        x2 = float(obj.find("bndbox/xmax").text)
        y2 = float(obj.find("bndbox/ymax").text)

        cx = ((x1 + x2) / 2) / img_w
        cy = ((y1 + y2) / 2) / img_h
        w = (x2 - x1) / img_w
        h = (y2 - y1) / img_h

        w_txt.write(f"{clas_index} {cx} {cy} {w} {h}" + "\n")

        # with open(txt_path, "a+") as f:
        #     f.write(f"{clas_index} {cx} {cy} {w} {h}" + "\n")
        #     # exit()

    w_txt.close()


# 解析labelme json
def labelme(json_path, txt_path):
    d = json.load(open(json_path))
    imgh = d["imageHeight"]
    imgw = d["imageWidth"]
    for target in d["shapes"]:
        label = target["label"]
        label_index = clas_name.index(label)  # 类别索引
        points = target["points"]

        x1, y1, x2, y2 = points[0][0], points[0][1], points[1][0], points[1][1]

        cx = ((x1 + x2) / 2) / imgw
        cy = ((y1 + y2) / 2) / imgh
        w = (x2 - x1) / imgw
        h = (y2 - y1) / imgh

        # 每个目标都要打开一次追加
        with open(txt_path, "a+") as f:
            f.write(f"{label_index} {cx} {cy} {w} {h}" + "\n")


if __name__ == '__main__':
    xmls_path = r".\data\annotations"
    save_dir = r".\data\labels"
    # images_path = r"\data\images"

    clas_name = ["yashang"]

    for xml_name in tqdm(os.listdir(xmls_path)):
        # print(xml_name)
        if xml_name[0] == ".":
            continue
        xml_path = os.path.join(xmls_path, xml_name)
        txt_path = os.path.join(save_dir, xml_name.split(".")[0] + ".txt")
        # print(xml_path)
        # print(txt_path)

        labelimg(xml_path, txt_path)
        # labelme(xml_path, txt_path)

    "********************************************************************************"

    # no_images_path = r"D:\my_program\project\L_dianchi-20230404\data\L_dianchi_1\all_data\crop\no_images"
    #
    # # 使没有目标框的图片有空的训练标签
    # for img_name in os.listdir(no_images_path):
    #     txt_path = os.path.join(txts_path, img_name.split(".")[0] + ".txt")
    #     if not os.path.exists(txt_path):
    #         with open(txt_path, "w") as f:
    #             continue
